import numpy as np
import pandas as pd
from dataclasses import dataclass
from typing import Dict, List, Optional
import numba
from scipy.linalg import inv
import time

@dataclass
class PlatParaInfo:
    """我方平台参数信息"""
    fLat: float  # 纬度 (度)
    fLon: float  # 经度 (度)
    fHeight: float  # 高度 (米)
    fSpeed: float  # 速度 (米/秒)
    fRoll: float  # 翻滚角 (度)
    fPitch: float  # 俯仰角 (度)
    fYaw: float  # 偏航角 (度)
    ID: int  # 平台ID
    timestamp: float  # 时间戳

@dataclass
class RadarDetection:
    """雷达探测信息"""
    fLat: float  # 纬度 (度)
    fLon: float  # 经度 (度)
    fHeight: float  # 高度 (米)
    timestamp: float  # 时间戳
    ID: int  # 目标ID

class HighPrecisionThreatAssessment:
    """高精度实时威胁评估系统"""
    
    def __init__(self):
        # WGS84椭球体参数
        self.a = 6378137.0  # 长半轴
        self.f = 1/298.257223563  # 扁率
        self.e2 = 2*self.f - self.f**2
        
        # 缓存数据结构
        self.ownship_history = {}
        self.target_history = {}
        self.kalman_filters = {}
        
        # 威胁评估参数
        self.threat_params = self._initialize_threat_params()
        
        # 性能优化参数
        self.max_history_length = 100
        self.min_time_interval = 0.01  # 最小时间间隔10ms
        
    def _initialize_threat_params(self):
        """初始化威胁评估参数"""
        return {
            # 距离威胁参数
            'k_r': 0.5, 'R_0': 10000, 'alpha_r': 0.1,
            # 角度威胁参数
            'sigma_psi': np.radians(30), 'sigma_theta': np.radians(20),
            'w_psi': 0.4, 'w_theta': 0.3, 'w_dpsi': 0.15, 'w_dtheta': 0.15,
            # 接近速度威胁参数
            'k_v': 0.3, 'V_crit': 200,
            # 机动威胁参数
            'A_max': 9.8 * 5, 'psi_max': np.radians(45), 'theta_max': np.radians(30),
            # 时间序列参数
            'beta1': 0.7, 'beta2': 0.3, 'delta_t_predict': 2.0
        }
    
    @staticmethod
    @numba.jit(nopython=True, fastmath=True)
    def llh_to_ecef(lat, lon, height):
        """LLH到ECEF坐标转换（优化版本）"""
        lat_rad = np.radians(lat)
        lon_rad = np.radians(lon)
        
        # WGS84参数
        a = 6378137.0
        e2 = 6.69437999014e-3
        
        N = a / np.sqrt(1 - e2 * np.sin(lat_rad)**2)
        
        x = (N + height) * np.cos(lat_rad) * np.cos(lon_rad)
        y = (N + height) * np.cos(lat_rad) * np.sin(lon_rad)
        z = (N * (1 - e2) + height) * np.sin(lat_rad)
        
        return x, y, z
    
    @staticmethod
    @numba.jit(nopython=True, fastmath=True)
    def ecef_to_enu_matrix(ref_lat, ref_lon):
        """计算ECEF到ENU的转换矩阵（优化版本）"""
        lat_rad = np.radians(ref_lat)
        lon_rad = np.radians(ref_lon)
        
        sin_lat = np.sin(lat_rad)
        cos_lat = np.cos(lat_rad)
        sin_lon = np.sin(lon_rad)
        cos_lon = np.cos(lon_rad)
        
        # ENU转换矩阵
        R = np.zeros((3, 3))
        R[0, 0] = -sin_lon
        R[0, 1] = cos_lon
        R[0, 2] = 0
        
        R[1, 0] = -sin_lat * cos_lon
        R[1, 1] = -sin_lat * sin_lon
        R[1, 2] = cos_lat
        
        R[2, 0] = cos_lat * cos_lon
        R[2, 1] = cos_lat * sin_lon
        R[2, 2] = sin_lat
        
        return R
    
    def calculate_relative_position(self, ownship: PlatParaInfo, target: RadarDetection):
        """计算相对位置（ENU坐标系）"""
        # 转换为ECEF坐标
        own_x, own_y, own_z = self.llh_to_ecef(ownship.fLat, ownship.fLon, ownship.fHeight)
        tar_x, tar_y, tar_z = self.llh_to_ecef(target.fLat, target.fLon, target.fHeight)
        
        # 计算相对位置
        delta_x = tar_x - own_x
        delta_y = tar_y - own_y
        delta_z = tar_z - own_z
        
        # 计算ENU转换矩阵
        R = self.ecef_to_enu_matrix(ownship.fLat, ownship.fLon)
        
        # 转换为ENU坐标
        enu = R @ np.array([delta_x, delta_y, delta_z])
        
        return enu
    
    def time_series_alignment(self, current_data, historical_data, target_time):
        """时间序列对齐"""
        if len(historical_data) < 2:
            return current_data
        
        # 找到最近的两个时间点
        times = np.array([data['timestamp'] for data in historical_data])
        idx = np.argsort(np.abs(times - target_time))[:2]
        
        if len(idx) < 2:
            return current_data
            
        data1, data2 = historical_data[idx[0]], historical_data[idx[1]]
        t1, t2 = times[idx[0]], times[idx[1]]
        
        if t1 == t2:
            return data1
        
        # 线性插值
        alpha = (target_time - t1) / (t2 - t1)
        aligned_data = {}
        
        for key in data1.keys():
            if isinstance(data1[key], (int, float)):
                aligned_data[key] = data1[key] + alpha * (data2[key] - data1[key])
            else:
                aligned_data[key] = data1[key]
                
        return aligned_data
    
    class AdaptiveKalmanFilter:
        """自适应卡尔曼滤波器"""
        def __init__(self, state_dim=10, measure_dim=3):
            self.state_dim = state_dim
            self.measure_dim = measure_dim
            
            # 状态向量: [E, N, U, V_E, V_N, V_U, psi, theta, dpsi, dtheta]
            self.x = np.zeros(state_dim)
            self.P = np.eye(state_dim) * 1000  # 初始协方差
            
            # 过程噪声和测量噪声
            self.Q = np.eye(state_dim) * 0.1
            self.R = np.eye(measure_dim) * 1.0
            
        def predict(self, dt):
            """预测步骤"""
            F = self._build_state_transition_matrix(dt)
            self.x = F @ self.x
            self.P = F @ self.P @ F.T + self.Q
            
        def update(self, z, H):
            """更新步骤"""
            y = z - H @ self.x  # 残差
            S = H @ self.P @ H.T + self.R  # 残差协方差
            K = self.P @ H.T @ inv(S)  # 卡尔曼增益
            
            self.x = self.x + K @ y
            self.P = (np.eye(self.state_dim) - K @ H) @ self.P
            
        def _build_state_transition_matrix(self, dt):
            """构建状态转移矩阵"""
            F = np.eye(self.state_dim)
            
            # 位置和速度关系
            F[0, 3] = dt
            F[1, 4] = dt
            F[2, 5] = dt
            
            # 角度和角速度关系
            F[6, 8] = dt
            F[7, 9] = dt
            
            return F
    
    def estimate_motion_parameters(self, target_id, current_enu, current_time):
        """估计运动参数"""
        if target_id not in self.target_history or len(self.target_history[target_id]) < 2:
            return None, None, None, None
        
        history = self.target_history[target_id]
        if len(history) < 2:
            return None, None, None, None
        
        # 获取最近两帧数据
        current_data = history[-1]
        prev_data = history[-2]
        
        dt = current_time - prev_data['timestamp']
        if dt <= 0:
            dt = self.min_time_interval
        
        # 计算方位角和俯仰角
        E, N, U = current_enu
        psi_current = np.arctan2(E, N)  # 方位角
        theta_current = np.arctan2(U, np.sqrt(E**2 + N**2))  # 俯仰角
        
        E_prev, N_prev, U_prev = prev_data['enu']
        psi_prev = np.arctan2(E_prev, N_prev)
        theta_prev = np.arctan2(U_prev, np.sqrt(E_prev**2 + N_prev**2))
        
        # 计算变化率
        dpsi = (psi_current - psi_prev) / dt
        dtheta = (theta_current - theta_prev) / dt
        
        return psi_current, theta_current, dpsi, dtheta
    
    def calculate_range_threat(self, R, dR):
        """计算距离威胁"""
        params = self.threat_params
        T_range = 1 / (1 + np.exp(params['k_r'] * (R / params['R_0'] - 1)))
        T_range *= (1 + params['alpha_r'] * abs(dR))
        
        return np.clip(T_range, 0, 1)
    
    def calculate_angle_threat(self, psi_rel, theta_rel, dpsi_rel, dtheta_rel):
        """计算角度威胁"""
        params = self.threat_params
        
        # 方位角威胁
        T_azimuth = np.exp(-(psi_rel**2) / (2 * params['sigma_psi']**2))
        
        # 俯仰角威胁
        T_elevation = np.exp(-(theta_rel**2) / (2 * params['sigma_theta']**2))
        
        # 综合角度威胁
        T_angle = (params['w_psi'] * T_azimuth + 
                  params['w_theta'] * T_elevation +
                  params['w_dpsi'] * abs(dpsi_rel) / params['psi_max'] +
                  params['w_dtheta'] * abs(dtheta_rel) / params['theta_max'])
        
        return np.clip(T_angle, 0, 1)
    
    def calculate_closure_threat(self, V_closure):
        """计算接近速度威胁"""
        params = self.threat_params
        T_closure = 1 / (1 + np.exp(-params['k_v'] * (V_closure / params['V_crit'] - 1)))
        return np.clip(T_closure, 0, 1)
    
    def calculate_maneuver_threat(self, acceleration, dpsi, dtheta, dt):
        """计算机动威胁"""
        params = self.threat_params
        
        # 线性机动威胁
        T_linear = np.linalg.norm(acceleration) / params['A_max']
        T_linear *= np.exp(-dt / 2.0)  # 时间衰减
        
        # 角机动威胁
        T_angular = (abs(dpsi) + abs(dtheta)) / (params['psi_max'] + params['theta_max'])
        
        return np.clip(T_linear, 0, 1), np.clip(T_angular, 0, 1)
    
    def update_threat_assessment(self, ownship: PlatParaInfo, radar_detections: List[RadarDetection]):
        """更新威胁评估"""
        threats = {}
        
        # 更新我方历史数据
        self._update_ownship_history(ownship)
        
        for detection in radar_detections:
            target_id = detection.ID
            
            # 计算相对位置
            enu = self.calculate_relative_position(ownship, detection)
            R = np.linalg.norm(enu)  # 距离
            
            # 更新目标历史
            self._update_target_history(target_id, enu, detection.timestamp)
            
            # 估计运动参数
            psi, theta, dpsi, dtheta = self.estimate_motion_parameters(
                target_id, enu, detection.timestamp)
            
            if psi is None:
                continue
            
            # 计算相对角度（考虑我方姿态）
            psi_rel = psi - np.radians(ownship.fYaw) - np.radians(ownship.fRoll) * np.sin(np.radians(ownship.fPitch))
            theta_rel = theta - np.radians(ownship.fPitch)
            
            # 计算接近速度
            if len(self.target_history[target_id]) >= 2:
                current_R = np.linalg.norm(self.target_history[target_id][-1]['enu'])
                prev_R = np.linalg.norm(self.target_history[target_id][-2]['enu'])
                dt = self.target_history[target_id][-1]['timestamp'] - self.target_history[target_id][-2]['timestamp']
                V_closure = -(current_R - prev_R) / dt if dt > 0 else 0
            else:
                V_closure = 0
            
            # 计算各威胁分量
            T_range = self.calculate_range_threat(R, V_closure)
            T_angle = self.calculate_angle_threat(psi_rel, theta_rel, dpsi, dtheta)
            T_closure = self.calculate_closure_threat(V_closure)
            
            # 计算机动威胁
            acceleration = self._estimate_acceleration(target_id)
            T_linear, T_angular = self.calculate_maneuver_threat(
                acceleration, dpsi, dtheta, self.min_time_interval)
            
            # 综合威胁计算
            T_total = self._calculate_total_threat(
                T_range, T_angle, T_closure, T_linear, T_angular, 
                detection.timestamp, target_id)
            
            threats[target_id] = {
                'total_threat': T_total,
                'components': {
                    'range_threat': T_range,
                    'angle_threat': T_angle,
                    'closure_threat': T_closure,
                    'linear_maneuver': T_linear,
                    'angular_maneuver': T_angular
                },
                'position': enu,
                'distance': R,
                'relative_angles': [np.degrees(psi_rel), np.degrees(theta_rel)]
            }
        
        return threats
    
    def _update_ownship_history(self, ownship: PlatParaInfo):
        """更新我方历史数据"""
        key = ownship.ID
        if key not in self.ownship_history:
            self.ownship_history[key] = []
        
        self.ownship_history[key].append({
            'timestamp': ownship.timestamp,
            'position': [ownship.fLat, ownship.fLon, ownship.fHeight],
            'attitude': [ownship.fRoll, ownship.fPitch, ownship.fYaw],
            'speed': ownship.fSpeed
        })
        
        # 限制历史数据长度
        if len(self.ownship_history[key]) > self.max_history_length:
            self.ownship_history[key] = self.ownship_history[key][-self.max_history_length:]
    
    def _update_target_history(self, target_id, enu, timestamp):
        """更新目标历史数据"""
        if target_id not in self.target_history:
            self.target_history[target_id] = []
        
        self.target_history[target_id].append({
            'timestamp': timestamp,
            'enu': enu.copy()
        })
        
        # 限制历史数据长度
        if len(self.target_history[target_id]) > self.max_history_length:
            self.target_history[target_id] = self.target_history[target_id][-self.max_history_length:]
    
    def _estimate_acceleration(self, target_id):
        """估计加速度"""
        history = self.target_history.get(target_id, [])
        if len(history) < 3:
            return np.zeros(3)
        
        # 使用差分法计算加速度
        current = history[-1]['enu']
        prev1 = history[-2]['enu']
        prev2 = history[-3]['enu']
        
        dt1 = history[-1]['timestamp'] - history[-2]['timestamp']
        dt2 = history[-2]['timestamp'] - history[-3]['timestamp']
        
        if dt1 <= 0 or dt2 <= 0:
            return np.zeros(3)
        
        velocity1 = (current - prev1) / dt1
        velocity2 = (prev1 - prev2) / dt2
        
        acceleration = (velocity1 - velocity2) / ((dt1 + dt2) / 2)
        
        return acceleration
    
    def _calculate_total_threat(self, T_range, T_angle, T_closure, T_linear, T_angular, timestamp, target_id):
        """计算综合威胁"""
        # 基础威胁分量
        T1 = T_range  # 距离威胁
        T2 = T_angle * (1 + 0.3 * T_angular)  # 角度威胁（考虑角机动）
        T3 = T_closure * (1 + 0.2 * T_linear)  # 接近速度威胁（考虑线性机动）
        
        # 威胁趋势分析
        T_trend = self._analyze_threat_trend(target_id, timestamp)
        T4 = T_trend if T_trend is not None else T_range
        
        # 权重调整（基于数据质量）
        weights = self._calculate_adaptive_weights(target_id)
        
        # 综合威胁
        T_total = (weights[0] * T1 + weights[1] * T2 + 
                  weights[2] * T3 + weights[3] * T4)
        
        # 交互威胁项
        T_interaction = 0.1 * np.sqrt(T1 * T2 + T2 * T3 + T1 * T3)
        
        return np.clip(T_total + T_interaction, 0, 1)
    
    def _analyze_threat_trend(self, target_id, current_time):
        """分析威胁趋势"""
        history = self.target_history.get(target_id, [])
        if len(history) < 3:
            return None
        
        # 简单的趋势分析（可扩展为更复杂的模型）
        recent_threats = []
        for i, data in enumerate(history[-3:]):
            # 这里简化计算，实际应该重新计算历史威胁
            R = np.linalg.norm(data['enu'])
            recent_threats.append(1 / (1 + np.exp(-0.001 * (10000 - R))))  # 简化威胁计算
        
        if len(recent_threats) >= 2:
            trend = recent_threats[-1] - recent_threats[0]
            return recent_threats[-1] + trend * 0.5  # 简单趋势预测
        
        return None
    
    def _calculate_adaptive_weights(self, target_id):
        """计算自适应权重"""
        # 基于数据稳定性的权重调整
        history = self.target_history.get(target_id, [])
        if len(history) < 2:
            return [0.25, 0.25, 0.25, 0.25]  # 默认等权重
        
        # 计算数据稳定性（位置变化的标准差）
        positions = np.array([data['enu'] for data in history])
        stability = 1.0 / (1.0 + np.std(positions, axis=0).mean())
        
        # 基于稳定性调整权重
        base_weights = [0.3, 0.25, 0.25, 0.2]  # 基础权重
        adjusted_weights = [w * stability for w in base_weights]
        
        # 归一化
        total = sum(adjusted_weights)
        return [w/total for w in adjusted_weights]

# 使用示例
def main():
    # 创建威胁评估系统
    threat_assessor = HighPrecisionThreatAssessment()
    
    for i in range(10000):
    # 模拟数据
        ownship = PlatParaInfo(
            fLat=39 + i * 0.001, fLon=116 + i * 0.001, fHeight=1000, fSpeed=200,
            fRoll=0, fPitch=0, fYaw=0, ID=1, timestamp=time.time()
        )
    
        radar_detections = [
            RadarDetection(fLat=41 - i * 0.001, fLon=118 -i *0.001, fHeight=1100, timestamp=time.time(), ID=1001),
            RadarDetection(fLat=42 - i*0.002, fLon=119 - i*0.002, fHeight=900, timestamp=time.time(), ID=1002)]
    
        # 进行威胁评估
        threats = threat_assessor.update_threat_assessment(ownship, radar_detections)
    
        # 输出结果
        for target_id, threat_info in threats.items():
            print(f"目标 {target_id}:")
            print(f"  综合威胁: {threat_info['total_threat']:.3f}")
            print(f"  距离: {threat_info['distance']:.1f}米")
            print(f"  威胁分量: {threat_info['components']}")

if __name__ == "__main__":
    main()